{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/force.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AbstractValue,\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    Box,\n",
    "    Capsule,\n",
    "    ContactResults,\n",
    "    DiagramBuilder,\n",
    "    InverseDynamicsController,\n",
    "    JointSliders,\n",
    "    LeafSystem,\n",
    "    MeshcatVisualizer,\n",
    "    MeshcatVisualizerParams,\n",
    "    MultibodyPlant,\n",
    "    Rgba,\n",
    "    RigidTransform,\n",
    "    RollPitchYaw,\n",
    "    Role,\n",
    "    RotationMatrix,\n",
    "    Simulator,\n",
    "    Sphere,\n",
    "    StartMeshcat,\n",
    "    StateInterpolatorWithDiscreteDerivative,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.scenarios import AddFloatingXyzJoint, AddShape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A simple writing example\n",
    "\n",
    "A popular example for class projects is to have a robot pick up a pen/pencil/chalk and write or draw something on the board. And that always begs the question \"how do we make the actual ink/graphite/chalk lines appear?\"\n",
    "\n",
    "Here is a simple teleop example with a LeafSystem that will help you visualize the drawing in meshcat.  You should be able to write on basically any (one) surface at any orientation.  Go ahead and try making the chalkboard into a sphere. ;-)\n",
    "\n",
    "You control the chalk, and see what you can draw.  \n",
    "- X : arrow left/right\n",
    "- Y : arrow down/up\n",
    "- Z : key A / Key D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Writing System.\n",
    "class MeshcatWriter(LeafSystem):\n",
    "    def __init__(\n",
    "        self,\n",
    "        meshcat,\n",
    "        drawing_geometry_ids,\n",
    "        canvas_geometry_ids,\n",
    "        rgba,\n",
    "        line_width=0.01,\n",
    "        draw_threshold=0.01,\n",
    "    ):\n",
    "        LeafSystem.__init__(self)\n",
    "        self._meshcat = meshcat\n",
    "        self._drawing_geometry_ids = drawing_geometry_ids\n",
    "        self._canvas_geometry_ids = canvas_geometry_ids\n",
    "\n",
    "        self._rgba = rgba\n",
    "        self._line_width = line_width\n",
    "        self._draw_threshold = draw_threshold\n",
    "\n",
    "        self.DeclareAbstractInputPort(\n",
    "            \"contact_results\", AbstractValue.Make(ContactResults())\n",
    "        )\n",
    "        self.DeclarePeriodicDiscreteUpdateEvent(0.01, 0, self.MaybeDraw)\n",
    "\n",
    "        self._p_WLastDraw_index = self.DeclareDiscreteState(3)\n",
    "        self._was_in_contact_index = self.DeclareDiscreteState([0])\n",
    "        self._num_drawn_index = self.DeclareDiscreteState([0])\n",
    "\n",
    "    def MaybeDraw(self, context, discrete_state):\n",
    "        results = self.get_input_port().Eval(context)\n",
    "\n",
    "        # The point on the canvas that most deeply penetrates the chalk.\n",
    "        p_WDraw = None\n",
    "        for i in range(results.num_hydroelastic_contacts()):\n",
    "            surface = results.hydroelastic_contact_info(i).contact_surface()\n",
    "            if (\n",
    "                surface.id_M() in self._drawing_geometry_ids\n",
    "                and surface.id_N() in self._canvas_geometry_ids\n",
    "            ):\n",
    "                p_WDraw = surface.centroid()\n",
    "                break\n",
    "            elif (\n",
    "                surface.id_N() in self._drawing_geometry_ids\n",
    "                and surface.id_M() in self._canvas_geometry_ids\n",
    "            ):\n",
    "                p_WDraw = surface.centroid()\n",
    "                break\n",
    "\n",
    "        if p_WDraw is not None:\n",
    "            p_WLastDraw = context.get_discrete_state(\n",
    "                self._p_WLastDraw_index\n",
    "            ).get_value()\n",
    "            was_in_contact = context.get_discrete_state(\n",
    "                self._was_in_contact_index\n",
    "            )[0]\n",
    "            num_drawn = context.get_discrete_state(self._num_drawn_index)[0]\n",
    "\n",
    "            length = np.linalg.norm(p_WDraw - p_WLastDraw)\n",
    "            if was_in_contact and length > self._draw_threshold:\n",
    "                meshcat.SetObject(\n",
    "                    f\"writer/{num_drawn}\",\n",
    "                    Capsule(self._line_width, length),\n",
    "                    self._rgba,\n",
    "                )\n",
    "                p_WMidpoint = (p_WDraw + p_WLastDraw) / 2\n",
    "                X_WCapsule = RigidTransform(\n",
    "                    RotationMatrix.MakeFromOneVector(p_WDraw - p_WLastDraw, 2),\n",
    "                    p_WMidpoint,\n",
    "                )\n",
    "                meshcat.SetTransform(f\"writer/{num_drawn}\", X_WCapsule)\n",
    "                discrete_state.set_value(\n",
    "                    self._num_drawn_index, [num_drawn + 1]\n",
    "                )\n",
    "                discrete_state.set_value(self._p_WLastDraw_index, p_WDraw)\n",
    "            elif not was_in_contact:\n",
    "                discrete_state.set_value(self._p_WLastDraw_index, p_WDraw)\n",
    "\n",
    "            discrete_state.set_value(self._was_in_contact_index, [1])\n",
    "        else:\n",
    "            discrete_state.set_value(self._was_in_contact_index, [0])\n",
    "\n",
    "\n",
    "def writing_example():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    time_step = 0.001\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step)\n",
    "    controller_plant = MultibodyPlant(time_step)\n",
    "    # Add the chalk to both\n",
    "    chalk_instance = AddShape(\n",
    "        plant,\n",
    "        Capsule(0.01, 0.2),\n",
    "        \"chalk\",\n",
    "        mass=1,\n",
    "        mu=1,\n",
    "        color=[1, 0.34, 0.2, 1.0],\n",
    "    )\n",
    "    AddFloatingXyzJoint(\n",
    "        plant, plant.GetFrameByName(\"chalk\"), chalk_instance, actuators=True\n",
    "    )\n",
    "    chalk_instance = AddShape(\n",
    "        controller_plant,\n",
    "        Capsule(0.01, 0.2),\n",
    "        \"chalk\",\n",
    "        mass=1,\n",
    "        mu=1,\n",
    "        color=[1, 0.34, 0.2, 1.0],\n",
    "    )\n",
    "    AddFloatingXyzJoint(\n",
    "        controller_plant,\n",
    "        controller_plant.GetFrameByName(\"chalk\"),\n",
    "        controller_plant.GetModelInstanceByName(\"chalk\"),\n",
    "        actuators=True,\n",
    "    )\n",
    "\n",
    "    # Add a writing surface to the sim plant only\n",
    "    board_instance = AddShape(\n",
    "        plant, Box(2, 2, 0.2), \"board\", color=[0.05, 0.05, 0.05, 1]\n",
    "    )\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"board\"),\n",
    "        RigidTransform([0, 0, -0.1]),\n",
    "    )\n",
    "\n",
    "    plant.Finalize()\n",
    "    controller_plant.Finalize()\n",
    "\n",
    "    q0 = [0, 0, 0.2]\n",
    "    plant.SetDefaultPositions(q0)\n",
    "\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "    meshcat.Delete()\n",
    "    meshcat.DeleteAddedControls()\n",
    "\n",
    "    # InverseDynamicsController\n",
    "    # TODO(russt): Tune the gains better\n",
    "    controller = builder.AddSystem(\n",
    "        InverseDynamicsController(\n",
    "            controller_plant,\n",
    "            kp=[20] * 3,\n",
    "            ki=[0] * 3,\n",
    "            kd=[10] * 3,\n",
    "            has_reference_acceleration=False,\n",
    "        )\n",
    "    )\n",
    "    builder.Connect(\n",
    "        plant.get_state_output_port(chalk_instance),\n",
    "        controller.get_input_port_estimated_state(),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        controller.get_output_port_control(), plant.get_actuation_input_port()\n",
    "    )\n",
    "\n",
    "    # Pose Sliders\n",
    "    teleop = builder.AddSystem(\n",
    "        JointSliders(\n",
    "            meshcat,\n",
    "            plant,\n",
    "            q0,\n",
    "            [-1, -1, -0.1],\n",
    "            [1, 1, 0.2],\n",
    "            0.02,\n",
    "            decrement_keycodes=[\"ArrowLeft\", \"ArrowDown\", \"KeyA\"],\n",
    "            increment_keycodes=[\"ArrowRight\", \"ArrowUp\", \"KeyD\"],\n",
    "        )\n",
    "    )\n",
    "\n",
    "    desired_state_from_position = builder.AddSystem(\n",
    "        StateInterpolatorWithDiscreteDerivative(\n",
    "            3, time_step, suppress_initial_transient=True\n",
    "        )\n",
    "    )\n",
    "    builder.Connect(\n",
    "        teleop.get_output_port(), desired_state_from_position.get_input_port()\n",
    "    )\n",
    "    builder.Connect(\n",
    "        desired_state_from_position.get_output_port(),\n",
    "        controller.get_input_port_desired_state(),\n",
    "    )\n",
    "\n",
    "    inspector = scene_graph.model_inspector()\n",
    "    chalk_frame_id = plant.GetBodyFrameIdOrThrow(\n",
    "        plant.GetBodyByName(\"chalk\").index()\n",
    "    )\n",
    "    board_frame_id = plant.GetBodyFrameIdOrThrow(\n",
    "        plant.GetBodyByName(\"board\").index()\n",
    "    )\n",
    "    writer = builder.AddSystem(\n",
    "        MeshcatWriter(\n",
    "            meshcat,\n",
    "            inspector.GetGeometries(chalk_frame_id, Role.kProximity),\n",
    "            inspector.GetGeometries(board_frame_id, Role.kProximity),\n",
    "            Rgba(1, 0.34, 0.2, 1.0),\n",
    "            line_width=0.005,\n",
    "        )\n",
    "    )\n",
    "    builder.Connect(\n",
    "        plant.get_contact_results_output_port(), writer.get_input_port()\n",
    "    )\n",
    "\n",
    "    # Simulate.\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    simulator = Simulator(diagram)\n",
    "\n",
    "    if running_as_notebook:\n",
    "        simulator.set_target_realtime_rate(1.0)\n",
    "        meshcat.AddButton(\"Stop Simulation\", \"Escape\")\n",
    "        print(\"Press Escape to stop the simulation\")\n",
    "        while meshcat.GetButtonClicks(\"Stop Simulation\") < 1:\n",
    "            simulator.AdvanceTo(simulator.get_context().get_time() + 2.0)\n",
    "        meshcat.DeleteButton(\"Stop Simulation\")\n",
    "    else:\n",
    "        simulator.AdvanceTo(0.1)\n",
    "\n",
    "\n",
    "writing_example()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "51 / 255"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Motion Planning.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.10.8 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
